################################################################################
# Lines.py - This module contains a generic Line object, along with a number of
#            subclasses that represent spectral lines common to Galaxy and QSO
#            spectra.
#

__author__ = 'Adrian Price-Whelan <adrn@astro.columbia.edu>'

# Standard library dependencies
import sys, os

# Third party dependencies
import numpy as np
import matplotlib.text as mtext

class Line:
    """ Represents a spectral line. """
    
    def __init__(self, wavelength, name, color='b', linestyle='-'):
        """ - wavelength must be numeric, and is ignorant of units
            - name must be a string
            - you may set the line's color and linestyle when creating the object, 
                all other matplotlib options must be passed as args or kwargs using the
                plot() method.
            - to show line labels, use 'show_labels=True' when using plot()
        
        """
        self.color = color
        self.linestyle = linestyle
        
        self.multi = False
        if isinstance(wavelength, [].__class__):
            # This line is a multi-plet or is split
            self.multi = True
            try:
                self.wavelength = map(float, wavelength)
            except ValueError:
                raise ValueError("Wavelength values must be numeric!")
        else:
            try:
                self.wavelength = float(wavelength)
            except ValueError:
                raise ValueError("Wavelength value must be numeric!")
        
        self.name = name
    
    def plot(self, axes, z=0.0, *args, **kwargs):
        if not kwargs.has_key("color") and not kwargs.has_key("c"):
            kwargs['color'] = self.color
        if not kwargs.has_key("linestyle") and not kwargs.has_key("ls"):
            kwargs['linestyle'] = self.linestyle
        
        if kwargs.has_key("show_labels") and kwargs["show_labels"]:
            show_labels = True
            del kwargs["show_labels"]
        else:
            show_labels = False
        
        if self.multi:
            for wv in self.wavelength: 
                line = axes.axvline(wv * (1. + z), label=self.name, *args, **kwargs)
                if show_labels:
                    a = axes.annotate(self.name, (wv * (1. + z), axes.get_ylim()[1]-0.1))
                    a.draggable()
        else:
            line = axes.axvline(self.wavelength * (1. + z), label=self.name, *args, **kwargs)
            if show_labels:
                a = axes.annotate(self.name, (self.wavelength * (1. + z), axes.get_ylim()[1]-0.1))
                a.draggable()

class LineGroup:
    """ Represents a group of spectral lines, e.g. the Balmer series, or 
        typical Galaxy Emission lines
    """
    
    def __init__(self, lines, color='b', linestyle='-'):
        """ - lines must be a list of Line() objects (see class above) 
            - color must either be a single character string, or a list of strings
                of length len(lines)
            - linestyle must either be a string, or a list of strings of length len(lines)
        """
        self.lines = lines
        
        if isinstance(color, [].__class__) and len(color) != len(lines):
            raise ValueError("Length of color list must be same as line list!")
        self.color = color
                
        if isinstance(linestyle, [].__class__) and len(linestyle) != len(lines):
            raise ValueError("Length of linestyle list must be same as line list!")
        self.linestyle = linestyle
        
    def plot(self, axes, z=0.0, *args, **kwargs):
        for ii, line in enumerate(self.lines):
            if not kwargs.has_key("color") and not kwargs.has_key("c"):
                if isinstance(self.color, [].__class__):
                    kwargs['color'] = self.color[ii]
                else:
                    kwargs['color'] = self.color
                    
            if not kwargs.has_key("linestyle") and not kwargs.has_key("ls"):
                if isinstance(self.linestyle, [].__class__):
                    kwargs['linestyle'] = self.linestyle[ii]
                else:
                    kwargs['linestyle'] = self.linestyle
            line.plot(axes, z=z, *args, **kwargs)

# Strong Galaxy Emission Lines
OI_6300 = Line(6300.23, "[OI] 6300", color='b', linestyle="--")
OII_3727 = Line([3726.1, 3728.8], "[OII] 3726, 3728", color='b', linestyle="--")
OIII_5007 = Line(5006.8, "[OIII] 5007", color='b', linestyle="--")
OIII_4959 = Line(4959.0, "[OIII] 4959", color='b', linestyle="--")
H_alpha = Line(6562.8, "H_alpha 6563", color='b', linestyle="--")
H_beta = Line(4861.3, "H_beta 4861", color='b', linestyle="--")
H_gamma = Line(4340.0, "H_gamma 4340", color='b', linestyle="--")
H_delta = Line(4102.8, "H_delta 4102", color='b', linestyle="--")
SII_6723 = Line([6716.0, 6731.0], "[SII] 6716, 6731", color='b', linestyle="--")
MgII_2798 = Line(2798.0, "MgII 2798", color='b', linestyle="--")
NII_6565 = Line([6548.0, 6583.0], "[NII] 6548, 6583", color='b', linestyle="--")
CIV_1549 = Line(1549.0, "CIV 1549", color='b', linestyle="--")

GalaxyEmissionLines = LineGroup([OI_6300,\
                                OII_3727,\
                                OIII_5007,\
                                OIII_4959,\
                                H_alpha,\
                                H_beta,\
                                H_gamma,\
                                H_delta,\
                                SII_6723,\
                                MgII_2798,\
                                NII_6565])

# Absorption Lines
CaII_3934 = Line(3933.7, "CaII (H) 3934", color='r', linestyle="--")
Ca_H = CaII_3934
CaII_3969 = Line(3968.5, "CaII (K) 3969", color='r', linestyle="--")
Ca_K = CaII_3969
G_band = Line(4304.4, "G_band 4304", color='r', linestyle="--")
Mg_5175 = Line(5175.3, "Mg 5175", color='r', linestyle="--")
Na_5894 = Line(5894.0, "Na 5894", color='r', linestyle="--")

GalaxyAbsorptionLines = LineGroup([CaII_3934,\
                                   CaII_3969,\
                                   G_band,\
                                   Mg_5175,\
                                   Na_5894])

# Series
BalmerLines = LineGroup([Line(6563.0, name="H_alpha"), \
                         Line(4861.0, name="H_beta"), \
                         Line(4341.0, name="H_gamma"), \
                         Line(4102.0, name="H_detla"), \
                         Line(3970.0, name="H_epsilon"), \
                         Line(3889.0, name="H_zeta"), \
                         Line(3835.0, name="H_eta")])

LymanLines = LineGroup([Line(1216.0, name="Ly_alpha"), \
                         Line(1026.0, name="Ly_beta"), \
                         Line(972.0, name="Ly_gamma"), \
                         Line(949.0, name="Ly_detla"), \
                         Line(937.0, name="Ly_epsilon"), \
                         Line(930.0, name="Ly_zeta"), \
                         Line(926.0, name="Ly_eta")])
        

if __name__ == "__main__":
    # Tests
    import matplotlib.pyplot as plt
    
    fig = plt.figure(figsize=(15,10))
    ax = fig.add_subplot(211)
    BalmerLines.plot(ax)
    LymanLines.plot(ax, color='g')
    
    ax2 = fig.add_subplot(212)
    GalaxyEmissionLines.plot(ax2, alpha=0.4)
    GalaxyAbsorptionLines.plot(ax2, alpha=0.4, show_labels=True)
    
    plt.show()
        

del sys, os, np, mtext